#include "stdafx.h"

#include <strsafe.h>
#include <Sddl.h>
#include <Userenv.h>
#include <AccCtrl.h>
#include <Aclapi.h>

#pragma comment(lib, "Userenv.lib")

//List of allowed capabilities for the application
extern WELL_KNOWN_SID_TYPE app_capabilities[] =
{
    WinCapabilityPrivateNetworkClientServerSid,
};

WCHAR container_name[] = L"MyTestAppContainer";
WCHAR container_desc[] = L"My AppContainer test sandbox";

BOOL IsInAppContainer();
BOOL SetSecurityCapabilities(PSID container_sid, SECURITY_CAPABILITIES *capabilities, PDWORD num_capabilities);

/*
    Create a container with container_name and run the specified application inside it
*/
BOOL RunExecutableInContainer(wchar_t *executable_path)
{
    PSID sid = NULL;
    HRESULT result;
    SECURITY_CAPABILITIES SecurityCapabilities = {0};
    DWORD num_capabilities = 0;
    SIZE_T attribute_size = 0;
    STARTUPINFOEXW startup_info = {0};
    PROCESS_INFORMATION process_info = {0};
    wchar_t desktop_file[MAX_PATH];
    HANDLE file_handle = INVALID_HANDLE_VALUE;
    wchar_t *string_sid = NULL;
    BOOL success = FALSE;
    wchar_t CommandLine[MAX_PATH] = { 0 };

    wcscpy_s(CommandLine, MAX_PATH, executable_path);
    startup_info.StartupInfo.cb = sizeof(STARTUPINFOEXW);

    do //Not a loop
    { 
        result = CreateAppContainerProfile(container_name, container_name, container_desc, NULL, 0, &sid);
        if(!SUCCEEDED(result))
        {
            if(HRESULT_CODE(result) == ERROR_ALREADY_EXISTS)
            {
                result = DeriveAppContainerSidFromAppContainerName(container_name, &sid);
                if(!SUCCEEDED(result))
                {
                    printf("Failed to get existing AppContainer name, error code: %d", HRESULT_CODE(result));
                    break;
                }
            }else{
                printf("Failed to create AppContainer, last error: %d\n", HRESULT_CODE(result));
                break;
            }   
        }

        printf("[Container Info]\nname: %ws\ndescription: %ws\n", container_name, container_desc);

        if(ConvertSidToStringSidW(sid, &string_sid))
            printf("Sid: %ws\n\n", string_sid);

        if(!SetSecurityCapabilities(sid, &SecurityCapabilities, &num_capabilities))
        {
            printf("Failed to set security capabilities, last error: %d\n", GetLastError());
            break;
        }

        ExpandEnvironmentStringsW(L"%userprofile%\\desktop\\allowed_test.txt", desktop_file, MAX_PATH-1);

        file_handle = CreateFileW(desktop_file, GENERIC_ALL, NULL, NULL, OPEN_ALWAYS, NULL, NULL);
        if(file_handle == INVALID_HANDLE_VALUE)
        {
            printf("Failed to create file %s, last error: %d\n", desktop_file, GetLastError());
            break;
        }
        
        if(!GrantNamedObjectAccess(sid, desktop_file, SE_FILE_OBJECT, FILE_ALL_ACCESS))
        {
            printf("Failed to grant explicit access to %s\n", desktop_file);
            break;
        }
        /*
        wchar_t* DllPath = (wchar_t*)L"C:\\Users\\Forrest\\Documents\\GitHub\\ExploitDev\\Experiments\\Tests\\AppContainer\\Release\\ScytheHookDll32.dll";
        if (!GrantNamedObjectAccess(sid, DllPath, SE_FILE_OBJECT, FILE_ALL_ACCESS))
        {
            printf("Failed to grant explicit access to %s\n", DllPath);
        }

        wchar_t* LogFilePath = (wchar_t*)L"C:\\ProgramData\\ScytheHook32.log";

        if (!GrantNamedObjectAccess(sid, LogFilePath, SE_FILE_OBJECT, FILE_ALL_ACCESS))
        {
            printf("Failed to grant explicit access to %ws\n", LogFilePath);
        }*/
        /*
        wchar_t* ObjectDirectory = (wchar_t*)L"\\BaseNamedObjects";

        if (!GrantNamedObjectAccess(sid, ObjectDirectory, SE_KERNEL_OBJECT, DIRECTORY_QUERY | DIRECTORY_TRAVERSE))
        {
            printf("Failed to grant explicit access to %ws\n", ObjectDirectory);
        }
        */
        InitializeProcThreadAttributeList(NULL, 1, NULL, &attribute_size);
        startup_info.lpAttributeList = (LPPROC_THREAD_ATTRIBUTE_LIST)malloc(attribute_size);

        if(!InitializeProcThreadAttributeList(startup_info.lpAttributeList, 1, NULL, &attribute_size))
        {
            printf("InitializeProcThreadAttributeList() failed, last error: %d", GetLastError());
            break;
        }

        if(!UpdateProcThreadAttribute(startup_info.lpAttributeList, 0, PROC_THREAD_ATTRIBUTE_SECURITY_CAPABILITIES, 
                                      &SecurityCapabilities, sizeof(SecurityCapabilities), NULL, NULL))
        {
            printf("UpdateProcThreadAttribute() failed, last error: %d", GetLastError());
            break;
        }

        if(!CreateProcessW(nullptr, CommandLine, NULL, NULL, FALSE, EXTENDED_STARTUPINFO_PRESENT, NULL, NULL,
                           (LPSTARTUPINFOW)&startup_info, &process_info))
        {
            printf("Failed to create process %s, last error: %d\n", executable_path, GetLastError());
            break;
        }

        printf("Successfully executed %s in AppContainer\n", executable_path);
        success = TRUE;

    } while (FALSE);

    if(startup_info.lpAttributeList)
        DeleteProcThreadAttributeList(startup_info.lpAttributeList);
 
    if(SecurityCapabilities.Capabilities)
        free(SecurityCapabilities.Capabilities);

    if(sid)
        FreeSid(sid);

    if(string_sid)
        LocalFree(string_sid);

    if(file_handle != INVALID_HANDLE_VALUE)
        CloseHandle(file_handle);

    if(file_handle != INVALID_HANDLE_VALUE && !success)
        DeleteFileW(desktop_file);

    return success;
}

/*
    Check if the current process is running inside an AppContainer
*/
BOOL IsInAppContainer()
{
    HANDLE process_token;
    BOOL is_container = 0; 
    DWORD return_length;

    OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &process_token);

    if (!GetTokenInformation(process_token, TokenIsAppContainer, &is_container, sizeof(is_container), &return_length)) 
        return false;

    return is_container;
}

/*
    Set the security capabilities of the container to those listed in app_capabilities
*/
BOOL SetSecurityCapabilities(PSID container_sid, SECURITY_CAPABILITIES *capabilities, PDWORD num_capabilities)
{
    DWORD sid_size = SECURITY_MAX_SID_SIZE;
    DWORD num_capabilities_ =  sizeof(app_capabilities) / sizeof(DWORD);
    SID_AND_ATTRIBUTES *attributes;
    BOOL success = TRUE;

    attributes = (SID_AND_ATTRIBUTES *)malloc(sizeof(SID_AND_ATTRIBUTES) * num_capabilities_);

    ZeroMemory(capabilities, sizeof(SECURITY_CAPABILITIES));
    ZeroMemory(attributes, sizeof(SID_AND_ATTRIBUTES) * num_capabilities_);

    for(unsigned int i = 0; i < num_capabilities_; i++)
    {
        attributes[i].Sid = malloc(SECURITY_MAX_SID_SIZE);
        if(!CreateWellKnownSid(app_capabilities[i], NULL, attributes[i].Sid, &sid_size))
        {
            success = FALSE;
            break;
        }
        attributes[i].Attributes = SE_GROUP_ENABLED;
    }

    if(success == FALSE)
    {
        for(unsigned int i = 0; i < num_capabilities_; i++)
        {
            if(attributes[i].Sid)
                LocalFree(attributes[i].Sid);
        }

        free(attributes);
        attributes = NULL;
        num_capabilities_ = 0;
    }

    capabilities->Capabilities = attributes;
    capabilities->CapabilityCount = num_capabilities_;
    capabilities->AppContainerSid = container_sid;
    *num_capabilities =  num_capabilities_;

    return success;
}
